from tqdm import tqdm

from yan import CvFo
import torch
import pandas as pd

voc = pd.read_pickle("voc_data.pandas_pickle")
net = CvFo(len(voc), 32, 8)
loss_func = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.0001)
data_set = pd.read_pickle("train_data.pandas_pickle")
batch_size = 12
bar = tqdm(len(100 * list(range(0, len(data_set), batch_size))))
for epoch in range(100):
    for i in range(0, len(data_set), batch_size):
        j = i + batch_size
        one_data = data_set[i:j]
        two_data = torch.Tensor(one_data).int()
        out = net(two_data[:, :-1])
        loss = loss_func(out.reshape([-1, out.shape[-1]]), two_data[:, 1:].reshape([-1]).long())
        bar.set_description("loss___{:.5f}".format(loss.item()))
        bar.update(1)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    if (epoch + 1) % 10 == 0:
        torch.save(net.state_dict(), "model_{}_loss_{:.5f}.pth".format(epoch + 1, loss.item()))
torch.save(net.state_dict(), "model_{}_loss_{:.5f}.pth".format(epoch + 1, loss.item()))
